﻿using SherpaOnnx;
using System;
using System.Collections.Generic;
using System.IO;
using System.Runtime.InteropServices;
using UnityEngine;
using UnityEngine.UI;

public class TTS : MonoBehaviour
{
    OfflineTtsConfig config;
    OfflineTts ot;
    OfflineTtsGeneratedAudio otga;
    OfflineTtsCallback callback;

    string Model;
    string Lexicon;
    string Tokens;
    string DataDir;
    string DictDir;
    string RuleFsts;
    public int SpeakerId = 1;
    public float Speed = 1f;

    int SampleRate = 22050;
    AudioSource audioSource;
    List<float> audioData = new List<float>();

    public Dropdown dpSpeakers;
    public InputField content;
    public InputField speakerid;
    public InputField speed;
    public Button generate;
    public AudioClip audioClip = null;
    /// <summary>
	/// 当前要读取的索引位置
	/// </summary>
	private int curAudioClipPos = 0;
    int modelIndex;

    public bool isInit = false;

    // Start is called before the first frame update
    void Start()
    {
        Loom.Initialize();
        audioSource = GetComponent<AudioSource>();
        string[] models = Enum.GetNames(typeof(Model));
        for (int i = 0; i < models.Length; i++)
        {
            Debug.Log(models[i]);
            dpSpeakers.options.Add(new Dropdown.OptionData(models[i]));
        }
        dpSpeakers.value = 0;
        dpSpeakers.captionText.text = models[0];
        OnDropDownValueChanged(0);
        dpSpeakers.onValueChanged.AddListener(OnDropDownValueChanged);

        generate.onClick.AddListener(() =>
        {
            audioSource.loop = false;
            audioSource.Stop();
            audioSource.clip = null;
            audioData.Clear();
            curAudioClipPos = 0;

            SpeakerId = Convert.ToInt32(speakerid.text);
            Speed = Convert.ToSingle(speed.text);
            if(!isInit)
            {
                Debug.LogWarning("未完成初始化");
                return;
            }
            Loom.RunAsync(() =>
            {
                Generate(content.text);
            });
        });
    }

    void OnDropDownValueChanged(int index)
    {
        Debug.Log("index:" + index);
        modelIndex = index;
        switch (index)
        {
            case 0:
                Model = "vits-zh-aishell3/vits-aishell3.onnx";
                Lexicon = "vits-zh-aishell3/lexicon.txt";
                Tokens = "vits-zh-aishell3/tokens.txt";
                //DataDir = "";
                //DictDir = "";
                RuleFsts = Application.streamingAssetsPath + "/vits-zh-aishell3/phone.fst" + ","
                + Application.streamingAssetsPath + "/vits-zh-aishell3/date.fst" + ","
                + Application.streamingAssetsPath + "/vits-zh-aishell3/number.fst";
                break;

            case 1:
                Model = "vits-zh-hf-theresa/theresa.onnx";
                Lexicon = "vits-zh-hf-theresa/lexicon.txt";
                Tokens = "vits-zh-hf-theresa/tokens.txt";
                //DataDir = "";
                DictDir = "vits-zh-hf-theresa/dict";
                RuleFsts = Application.streamingAssetsPath + "/vits-zh-hf-theresa/phone.fst" + ","
                + Application.streamingAssetsPath + "/vits-zh-hf-theresa/date.fst" + ","
                + Application.streamingAssetsPath + "/vits-zh-hf-theresa/number.fst";
                break;

            case 2:
                //666和99
                Model = "vits-zh-hf-eula/eula.onnx";
                Lexicon = "vits-zh-hf-eula/lexicon.txt";
                Tokens = "vits-zh-hf-eula/tokens.txt";
                //DataDir = "";
                DictDir = "vits-zh-hf-eula/dict";
                RuleFsts = Application.streamingAssetsPath + "/vits-zh-hf-eula/phone.fst" + ","
                + Application.streamingAssetsPath + "/vits-zh-hf-eula/date.fst" + ","
                + Application.streamingAssetsPath + "/vits-zh-hf-eula/number.fst";
                break;

            case 3:
                Model = "vits-zh-hf-fanchen-C/vits-zh-hf-fanchen-C.onnx";
                Lexicon = "vits-zh-hf-fanchen-C/lexicon.txt";
                Tokens = "vits-zh-hf-fanchen-C/tokens.txt";
                //DataDir = "";
                DictDir = "vits-zh-hf-fanchen-C/dict";
                RuleFsts = Application.streamingAssetsPath + "/vits-zh-hf-fanchen-C/phone.fst" + ","
                + Application.streamingAssetsPath + "/vits-zh-hf-fanchen-C/date.fst" + ","
                + Application.streamingAssetsPath + "/vits-zh-hf-fanchen-C/number.fst";
                break;

            case 4:
                Model = "vits-melo-tts-zh_en/model.onnx";
                Lexicon = "vits-melo-tts-zh_en/lexicon.txt";
                Tokens = "vits-melo-tts-zh_en/tokens.txt";
                //DataDir = "";
                DictDir = "vits-melo-tts-zh_en/dict";
                RuleFsts = Application.streamingAssetsPath + "/vits-melo-tts-zh_en/phone.fst" + ","
                    + Application.streamingAssetsPath + "/vits-melo-tts-zh_en/date.fst" + ","
                + Application.streamingAssetsPath + "/vits-melo-tts-zh_en/number.fst";
                break;
        }
        Loom.RunAsync(() =>
        {
            Init();
        });
    }

    void Init()
    {
        isInit = false;
        Debug.LogWarning("初始化改为异步，大概5秒左右，如果点击按钮太快，可能没反应？");
        if (ot != null)
        {
            ot.Dispose();
        }
        if (otga != null)
        {
            otga.Dispose();
        }
        config = new OfflineTtsConfig();
        config.Model.Vits.Model = Path.Combine(Application.streamingAssetsPath, Model);
        config.Model.Vits.Lexicon = Path.Combine(Application.streamingAssetsPath, Lexicon);
        config.Model.Vits.Tokens = Path.Combine(Application.streamingAssetsPath, Tokens);
        //config.Model.Vits.DataDir = Path.Combine(Application.streamingAssetsPath, DataDir);

        //vits-zh-aishell3没有dict
        if (modelIndex != 0)
        {
            config.Model.Vits.DictDir = Path.Combine(Application.streamingAssetsPath, DictDir);
        }
        config.Model.Vits.NoiseScale = 0.667f;
        config.Model.Vits.NoiseScaleW = 0.8f;
        config.Model.Vits.LengthScale = 1f;
        config.Model.NumThreads = 12;
        config.Model.Debug = 1;
        //支持cuda但C#不知道怎么用
        config.Model.Provider = "cpu";
        config.RuleFsts = RuleFsts;
        config.MaxNumSentences = 1;
        ot = new OfflineTts(config);
        SampleRate = ot.SampleRate;
        callback = new OfflineTtsCallback(MyCallback);
        isInit = true;
    }

    void Generate(string content)
    {
        otga = ot.GenerateWithCallback(content, Speed, SpeakerId, callback);
        string fileName = DateTime.Now.ToFileTime().ToString() + ".wav";
        bool ok = otga.SaveToWaveFile(Application.streamingAssetsPath + "/" + fileName);
        if (ok)
        {
            Debug.Log(fileName + " Save succeeded!");
        }
        else
        {
            Debug.Log("Failed");
        }
    }

    int MyCallback(IntPtr samples, int n)
    {
        float[] tempData = new float[n];
        Marshal.Copy(samples, tempData, 0, n);
        audioData.AddRange(tempData);
        Debug.Log("n:" + n);
        Loom.QueueOnMainThread(() =>
        {
            if (!audioSource.isPlaying && audioData.Count > SampleRate * 2)
            {
                audioClip = AudioClip.Create("SynthesizedAudio", SampleRate * 2, 1,
                    SampleRate, true, (float[] data) =>
                {
                    ExtractAudioData(data);
                });
                audioSource.clip = audioClip;
                audioSource.loop = true;
                audioSource.Play();
            }
        });
        return n;
    }

    bool ExtractAudioData(float[] data)
    {
        if (data == null || data.Length == 0)
        {
            return false;
        }
        bool hasData = false;//是否真的读取到数据
        int dataIndex = 0;//当前要写入的索引位置
        if (audioData != null && audioData.Count > 0)
        {
            while (curAudioClipPos < audioData.Count && dataIndex < data.Length)
            {
                data[dataIndex] = audioData[curAudioClipPos];
                curAudioClipPos++;
                dataIndex++;
                hasData = true;
            }
        }

        //剩余部分填0
        while (dataIndex < data.Length)
        {
            data[dataIndex] = 0;
            dataIndex++;
        }
        return hasData;
    }

    private void OnApplicationQuit()
    {
        if (ot != null)
        {
            ot.Dispose();
        }
        if (otga != null)
        {
            otga.Dispose();
        }
    }
}